"""
MDC（Mapped Diagnostic Context）上下文管理器

提供线程/协程安全的诊断上下文管理，用于在日志中跟踪请求级别的上下文信息
"""

import contextvars
import threading
import uuid
from contextlib import contextmanager
from datetime import datetime
from typing import Any, Dict, Optional

# 全局 MDC 上下文变量
_mdc_context: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar("mdc_context", default={})


class MDCContext:
    """MDC 上下文管理器类"""

    def __init__(self):
        self._local = threading.local()

    def put(self, key: str, value: Any) -> None:
        """向当前上下文添加键值对"""
        context = _mdc_context.get()
        context[key] = value
        _mdc_context.set(context)

    def get(self, key: str, default: Any = None) -> Any:
        """从当前上下文获取值"""
        context = _mdc_context.get()
        return context.get(key, default)

    def remove(self, key: str) -> None:
        """从当前上下文移除键值对"""
        context = _mdc_context.get()
        context.pop(key, None)
        _mdc_context.set(context)

    def clear(self) -> None:
        """清空当前上下文"""
        _mdc_context.set({})

    def get_context(self) -> Dict[str, Any]:
        """获取完整的上下文字典"""
        return _mdc_context.get().copy()

    def set_context(self, context: Dict[str, Any]) -> None:
        """设置完整的上下文"""
        _mdc_context.set(context.copy())


# 全局 MDC 实例
mdc = MDCContext()


@contextmanager
def mdc_context(**kwargs):
    """
    MDC 上下文管理器

    用法：
    with mdc_context(request_id="123", user_id=456):
        # 在此范围内的日志会自动包含上下文信息
        logger.info("处理请求")

    Args:
        **kwargs: 要设置的上下文键值对
    """
    # 保存当前上下文
    original_context = mdc.get_context()

    try:
        # 设置新的上下文
        new_context = original_context.copy()
        new_context.update(kwargs)
        mdc.set_context(new_context)
        yield
    finally:
        # 恢复原始上下文
        mdc.set_context(original_context)


@contextmanager
def request_mdc_context(request_id: Optional[str] = None, **kwargs):
    """
    请求级别的 MDC 上下文管理器

    自动生成请求ID并设置基础上下文信息

    Args:
        request_id: 请求ID，如果未提供则自动生成
        **kwargs: 额外的上下文信息
    """
    # 生成请求ID
    if not request_id:
        request_id = str(uuid.uuid4())

    # 基础上下文信息
    base_context = {"request_id": request_id, "start_time": datetime.now(), **kwargs}

    with mdc_context(**base_context):
        try:
            yield request_id
        finally:
            # 记录请求结束时间
            mdc.put("end_time", datetime.now())


def get_request_id() -> Optional[str]:
    """获取当前请求ID"""
    return mdc.get("request_id")


def get_user_id() -> Optional[int]:
    """获取当前用户ID"""
    return mdc.get("user_id")


def get_trace_id() -> Optional[str]:
    """获取跟踪ID（可用于分布式追踪）"""
    return mdc.get("trace_id") or mdc.get("request_id")
